package com.idanchuang.component.mq.amqp.pressure;

import com.alibaba.fastjson.JSON;
import com.idanchuang.component.core.helper.EnvContext;
import com.idanchuang.component.logback.helper.LogContext;
import com.idanchuang.component.mq.amqp.processor.AfterReceivePostProcessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.config.AbstractRabbitListenerContainerFactory;

import java.net.URLDecoder;
import java.util.Map;

/**
 * @author yjy
 * Created at 2021/11/9 3:26 下午
 */
public class PressureAfterReceivePostProcessor implements AfterReceivePostProcessor {

    private static final Logger abstractRabbitListenerContainerFactoryLogger = LoggerFactory.getLogger(AbstractRabbitListenerContainerFactory.class);

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        try {
            Map<String, Object> header = message.getMessageProperties().getHeaders();
            Object attributes = header.get(EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME);
            abstractRabbitListenerContainerFactoryLogger.debug("AfterReceivePostProcessors header: {}", header);
            // 清除当前环境变量
            EnvContext.clearContext();
            // 压测标
            String headerName = EnvContext.getEnvHeaderName();
            Object headerValue = header.get(headerName);
            if (headerValue != null) {
                abstractRabbitListenerContainerFactoryLogger.debug("AfterReceivePostProcessors get header: {}, value: {}", headerName, headerValue);
            }
            EnvContext.setCurrentEnv(headerValue != null ? headerValue.toString() : null);
            // 自定义链路id
            String dcTraceIdName = EnvContext.getDcTraceIdHeaderName();
            Object dcTraceIdValue = header.get(dcTraceIdName);
            if (dcTraceIdValue != null) {
                abstractRabbitListenerContainerFactoryLogger.debug("AfterReceivePostProcessors get header: {}, value: {}", dcTraceIdName, dcTraceIdValue);
            }
            if (attributes != null) {
                String attributesStr = URLDecoder.decode(attributes.toString(), "UTF-8");
                EnvContext.setAllAttribute(JSON.parseObject(attributesStr));
                abstractRabbitListenerContainerFactoryLogger.debug("AfterReceivePostProcessors get header: {}, value: {}", EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributesStr);
            }
            EnvContext.setDcTraceIdIfAbsent();
        } catch (Exception e) {
            // 降级
            LogContext.setKeyword("componentError");
            abstractRabbitListenerContainerFactoryLogger.error("Amqp setAfterReceivePostProcessors error", e);
        }
        return message;
    }
}
